import argparse
import openpyxl
import random
from collections import defaultdict

# 设置不同的种子，例如使用当前时间
random.seed()

def parse_arguments():
    parser = argparse.ArgumentParser(description="Process box data.")
    parser.add_argument("--input",  default='./box_data.xlsx', help="Path to the input Excel file")
    return parser.parse_args()



def process_box_data(input_file, num_wooden, num_pallet, num_cable, num_cardboard):
    # 打开工作簿
    workbook = openpyxl.load_workbook(input_file)

    # 选择要读取的工作表
    sheet = workbook.active  # 或者可以通过名称选择：sheet = workbook['Sheet1']

    # 创建一个空字典，用于存放符合条件的列表
    result_dict = defaultdict(list)

    # 创建包含关键词的字典
    keyword_dict = {'木箱', '托盘', '缆盘', '纸箱'}
    index_keyword_dict = {'木箱': 1, '托盘': 2, '缆盘': 0, '纸箱': 3}

    # 读取数据
    for row in sheet.iter_rows(min_row=1, max_row=sheet.max_row, values_only=True):
        if row[1] in keyword_dict:
            result_dict[row[1]].append([row[3], row[4], row[5], row[8]])

    # 输出结果字典
    result_dict = {key: set(tuple(sublist) for sublist in value) for key, value in result_dict.items()}

    ratios = {
        '木箱': num_wooden,
        '托盘': num_pallet,
        '缆盘': num_cable,
        '纸箱': num_cardboard
    }

    output_data = []

    for box_type, box_number in ratios.items():
        selected_boxes = random.choices(list(result_dict[box_type]), k=box_number)
        output_data.extend([[box[0], box[1], box[2], box[3], index_keyword_dict[box_type]] for box in selected_boxes])

    # 使用字典追踪每个子列表的出现次数
    count_dict = defaultdict(int)
    for sublist in output_data:
        tuple_sublist = tuple(sublist)
        count_dict[tuple_sublist] += 1

    # 调整维度并添加出现次数
    result = [[item[0], item[1], item[2], item[3], count, item[4]] for item, count in count_dict.items()]

    return result

# 随机生成货物数量
def data_generator():
    choices = [
        (range(1, 301), 0.968),
        (range(301, 601), 0.0131),
        (range(601, 901), 0.0039),
        (range(901, 1001), 0.0045),
        (range(1201, 1500), 0.0017),
        (range(1501, 1800), 0.0009),
        (range(1801, 2100), 0.0018),
        (range(2101, 2400), 0.0009),
        (range(2401, 2700), 0.0005),
        (range(2701, 3000), 0.0010),
        (range(3001, 3300), 0.0004),
        (range(3301, 3900), 0.0004),
        (range(3901, 4200), 0.0003),
        (range(4201, 4500), 0.0004),
        (range(4501, 4800), 0.0001),
        (range(4801, 5100), 0.0001),
        (range(5101, 5400), 0.0001),
        (range(5401, 5700), 0.0003),
        (range(5701, 6000), 0.0001),
        (range(6001, 6300), 0.0001),
        (range(6301, 6600), 0.0001),
        (range(6601, 6900), 0.0001),
        (range(6901, 7200), 0.0001),
        (range(7201, 7500), 0.0004),
        (range(11101, 11400), 0.0001),
        (range(11701, 12000), 0.0001),
        (range(12301, 12600), 0.0001)
    ]

    selected_range = random.choices(*zip(*choices))[0]
    return random.randint(selected_range.start, selected_range.stop - 1)

def split_with_adjustment(num_data, prob_wooden, prob_pallet, prob_cable, prob_cardboard):
    # 计算每个部分的大小
    # parts = [ratio * num_data for ratio in ratios]
    number_wooden = num_data * prob_wooden
    number_pallet = num_data * prob_pallet
    number_cable = num_data * prob_cable
    number_cardboard = num_data * prob_cardboard
    parts = [number_wooden, number_pallet, number_cable, number_cardboard]
    # 将浮点数部分转换为整数
    parts = [int(part) for part in parts]
    # 计算拆分后的总和和误差
    num_data_after_split = sum(parts)
    error = num_data - num_data_after_split

    # 如果有误差，将其分配给最大的部分
    if error != 0:
        max_part_index = parts.index(max(parts))
        parts[max_part_index] += error

    num_wooden = parts[0]
    num_pallet = parts[1]
    num_cable = parts[2]
    num_cardboard = parts[3]

    return num_wooden, num_pallet, num_cable, num_cardboard



def generate_box_data():
    args = parse_arguments()
    random_data = data_generator()
    if 1 < int(random_data) < 300:
        num_data = int(random_data)
        prob_wooden = 0.15
        prob_pallet = 0.11
        prob_cable = 0.01
        prob_cardboard = 0.73
    elif 301 < int(random_data) < 600:
        num_data = int(random_data)
        prob_wooden = 0.02
        prob_pallet = 0.05
        prob_cable = 0.01
        prob_cardboard = 0.91
    else:
        num_data = int(random_data)
        prob_wooden = 0
        prob_pallet = 0
        prob_cable = 0
        prob_cardboard = 1

    # 按比例拆分数据
    num_wooden, num_pallet, num_cable, num_cardboard = split_with_adjustment(num_data, prob_wooden, prob_pallet, prob_cable, prob_cardboard)

    # 按比例生成数据
    result = process_box_data(args.input, num_wooden, num_pallet, num_cable, num_cardboard)
    return result



if __name__ == "__main__":
    result = generate_box_data()
    print(result)










